import numpy as np
import mindspore as ms
import matplotlib.pyplot as plt
import cv2

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))

image = cv2.imread('images/truck.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

from segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "sam_vit_b_01ec64.ckpt"
model_type = "vit_b"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)

predictor = SamPredictor(sam)

predictor.set_image(image)
# torch_image_feature = np.load('truck_features.npy')
# predictor.set_torch_image(torch_image_feature)

# 1.单点输出
input_point = np.array([[500, 375]])
input_label = np.array([1])

masks, scores, logits_single = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)

# np.save("ms_logits_1.npy", logits_single)

# for i, (mask, score) in enumerate(zip(masks, scores)):
#     plt.figure(figsize=(10,10))
#     plt.imshow(image)
#     show_mask(mask, plt.gca())
#     show_points(input_point, input_label, plt.gca())
#     plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
#     plt.axis('off')
#     plt.show()

# 2.两点确定一个物体
input_point = np.array([[500, 375], [1125, 625]])
input_label = np.array([1, 1])

mask_input = logits_single[np.argmax(scores), :, :]  # Choose the model's best mask

masks, _, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    mask_input=None,#mask_input[None, :, :],
    multimask_output=False,
)

np.save("ms_logits_2.npy", logits)

plt.figure(figsize=(10,10))
plt.imshow(image)
show_mask(masks, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show() 

# 2.1 两点确定两个物体
input_point = ms.Tensor(np.array([
    [[500, 375]],
    [[1125, 625]],
]), ms.int32)
input_label = ms.Tensor(np.array([[1], [1]]), ms.int32)
transformed_points = predictor.transform.apply_coords_batch(input_point, image.shape[:2])
masks, _, logits = predictor.predict_tensor(
    point_coords=transformed_points,
    point_labels=input_label,
    multimask_output=False,
)
masks = masks.asnumpy()

np.save("ms_logits_3.npy", logits.asnumpy())

plt.figure(figsize=(10, 10))
plt.imshow(image)
for mask in masks:
    show_mask(mask, plt.gca(), random_color=True)
show_points(input_point.view(-1, 2).asnumpy(), input_label.view(-1).asnumpy(), plt.gca())
plt.axis('off')
plt.show()

# 3.输入一个前景点和背景点
input_point = np.array([[500, 375], [1125, 625]])
input_label = np.array([1, 0])

mask_input = logits_single[np.argmax(scores), :, :]  # Choose the model's best mask

masks, _, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    mask_input=mask_input[None, :, :],
    multimask_output=False,
)

np.save("ms_logits_4.npy", logits)

plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show()

# 4.根据坐标框确定mask
input_box = np.array([425, 600, 700, 875])
masks, _, logits = predictor.predict(
    point_coords=None,
    point_labels=None,
    box=input_box[None, :],
    multimask_output=False,
)

np.save("ms_logits_5.npy", logits)

plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[0], plt.gca())
show_box(input_box, plt.gca())
plt.axis('off')
plt.show()

# 5.联合坐标框和点
input_box = np.array([425, 600, 700, 875])
input_point = np.array([[575, 750]])
input_label = np.array([0])

masks, _, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    box=input_box,
    multimask_output=False,
)

np.save("ms_logits_6.npy", logits)

plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[0], plt.gca())
show_box(input_box, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show()

# 5.1 联合多个坐标框和多个点，确定不同物体
input_boxes = ms.Tensor(np.array([[425, 600, 700, 875], 
                                  [1360, 525, 1680, 780]]))
input_points = ms.Tensor(np.array([[[575, 750]], 
                                  [[1525, 670]]]))
input_labels = ms.Tensor(np.array([[1], [1]]))
transformed_boxes = predictor.transform.apply_boxes_batch(input_boxes, image.shape[:2])
transformed_points = predictor.transform.apply_coords_batch(input_points, image.shape[:2])
masks, _, logits = predictor.predict_tensor(
    point_coords=input_points,
    point_labels=input_labels,
    boxes=transformed_boxes,
    multimask_output=False,
)

np.save("ms_logits_7.npy", logits.asnumpy())

plt.figure(figsize=(10, 10))
plt.imshow(image)
for mask in masks:
    show_mask(mask.asnumpy(), plt.gca(), random_color=True)
for box in input_boxes:
    show_box(box.asnumpy(), plt.gca())
for point, label in zip(input_points, input_labels):
    show_points(point.asnumpy(), label.asnumpy(), plt.gca())
plt.axis('off')
plt.show()

# 6.多个框确定多个mask
input_boxes = ms.Tensor(np.array([
    [75, 275, 1725, 850],
    [425, 600, 700, 875],
    [1375, 550, 1650, 800],
    [1240, 675, 1400, 750],
]))
transformed_boxes = predictor.transform.apply_boxes_batch(input_boxes, image.shape[:2])
masks, _, logits = predictor.predict_tensor(
    point_coords=None,
    point_labels=None,
    boxes=transformed_boxes,
    multimask_output=False,
)

np.save("ms_logits_8.npy", logits.asnumpy())

plt.figure(figsize=(10, 10))
plt.imshow(image)
for mask in masks:
    show_mask(mask.asnumpy(), plt.gca(), random_color=True)
for box in input_boxes:
    show_box(box.asnumpy(), plt.gca())
plt.axis('off')
plt.show()

# 7.多图多坐标框
# image1 = image  # truck.jpg from above
# image1_boxes = ms.Tensor([
#     [75, 275, 1725, 850],
#     [425, 600, 700, 875],
#     [1375, 550, 1650, 800],
#     [1240, 675, 1400, 750],
# ])

# image2 = cv2.imread('images/groceries.jpg')
# image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
# image2_boxes = ms.Tensor(np.array([
#     [450, 170, 520, 350],
#     [350, 190, 450, 350],
#     [500, 170, 580, 350],
#     [580, 170, 640, 350],
# ]))

# from segment_anything.utils.transforms import ResizeLongestSide

# resize_transform = ResizeLongestSide(sam.image_encoder.img_size)

# def prepare_image(image, transform):
#     image = transform.apply_image(image)
#     image = ms.Tensor(image)
#     return image.transpose(2, 0, 1)

# batched_input = [
#      {
#          'image': prepare_image(image1, resize_transform),
#          'boxes': resize_transform.apply_boxes_batch(image1_boxes, image1.shape[:2]),
#          'original_size': image1.shape[:2]
#      },
#      {
#          'image': prepare_image(image2, resize_transform),
#          'boxes': resize_transform.apply_boxes_batch(image2_boxes, image2.shape[:2]),
#          'original_size': image2.shape[:2]
#      }
# ]
# batched_output = sam(batched_input, multimask_output=False)

# fig, ax = plt.subplots(1, 2, figsize=(20, 20))

# ax[0].imshow(image1)
# for mask in batched_output[0]['masks']:
#     show_mask(mask.asnumpy(), ax[0], random_color=True)
# for box in image1_boxes:
#     show_box(box.asnumpy(), ax[0])
# ax[0].axis('off')

# ax[1].imshow(image2)
# for mask in batched_output[1]['masks']:
#     show_mask(mask.asnumpy(), ax[1], random_color=True)
# for box in image2_boxes:
#     show_box(box.asnumpy(), ax[1])
# ax[1].axis('off')

# plt.tight_layout()
# plt.show()
